import pandas as pd

# ==================== 配置参数 ====================
INPUT_FILE = "ori_message.xlsx"  # 原始数据文件路径（可替换为绝对路径）
OUTPUT_FILE = "monthly_average_data.xlsx"  # 输出文件路径


# ==================== 数据读取与预处理 ====================
def process_data():
    # 读取Excel文件
    excel_file = pd.ExcelFile(INPUT_FILE)
    years = ['2016', '2017', '2018', '2019', '2020', '2021']  # 对应工作表名称
    all_monthly_data = []  # 存储所有年份的处理结果

    for year in years:
        # 读取当前年份的工作表
        df = excel_file.parse(year)

        # 处理缺失值（向前填充年、月）
        df['年'] = df['年'].ffill().astype(int)  # 填充并转为整数
        df['月'] = df['月'].ffill().astype(int)  # 填充并转为整数

        # 按年、月分组计算平均值（注意含沙量列名末尾有空格）
        grouped = df.groupby(['年', '月']).agg({
            '流量(m3/s)': 'mean',
            '含沙量(kg/m3) ': 'mean'  # 严格匹配原始列名（含末尾空格）
        }).reset_index()

        # 计算排沙量（流量 × 含沙量）
        grouped['排沙量(kg/s)'] = grouped['流量(m3/s)'] * grouped['含沙量(kg/m3) ']

        # 保留4位小数（提升可读性）
        grouped = grouped.round({
            '流量(m3/s)': 2,
            '含沙量(kg/m3) ': 4,
            '排沙量(kg/s)': 2
        })

        all_monthly_data.append(grouped)

    # 合并所有年份数据并按年、月排序
    result = pd.concat(all_monthly_data, ignore_index=True)
    result = result.sort_values(by=['年', '月']).reset_index(drop=True)

    # 重命名列（去除含沙量列名末尾的空格，可选）
    result = result.rename(columns={'含沙量(kg/m3) ': '含沙量(kg/m3)'})

    # 保存结果到Excel（可选）
    result.to_excel(OUTPUT_FILE, index=False)

    return result


# ==================== 输出处理后的数据 ====================
if __name__ == "__main__":
    # 执行数据处理
    processed_data = process_data()

    # 打印处理后的结果（表格形式）
    print("\n========== 每月平均流量、含沙量及排沙量 ==========\n")
    print(processed_data.to_string(index=False))  # 不显示索引列
